# coding=utf-8

from flask import session
from requests.cookies import RequestsCookieJar
from typing import Optional, List, Mapping

from app.extensions import session_id_manager
from app.cores.dictionaries import DISPATCHER_TYPE, DISPATCHER_TRIGGER_TYPE

# requests库会在发起请求时根据请求服务器的domain选择合适的cookie作为请求头中的cookie数据发出
#   可参考代码 requests.models.prepare_cookies()


class HTTPCookiePoolManager:
    """
    管理自动化测试中请求头Cookie数据

    CookiePool = {
        session_id1: {
            DISPATCHER_TYPE.DEBUG: RequestsCookieJar(),
            DISPATCHER_TYPE.BUILD: RequestsCookieJar(),
        },
        session_id2: {
            DISPATCHER_TYPE.DEBUG: RequestsCookieJar(),
            DISPATCHER_TYPE.BUILD: RequestsCookieJar(),
        },
        ...
    }

    ScheduleCookiePool = {
        project_id1: RequestsCookieJar(),
        project_id2: RequestsCookieJar(),
        ...
    }
    """

    # cookie池，不同的会话放到不同的字典中，key为区分会话唯一编号放入cookie中
    CookiePool = {}

    # 对于定时构建的项目，使用另一个Cookie池保存，key为项目id
    ScheduleCookiePool = {}

    @classmethod
    def update_cookie_pool(cls, rcj: RequestsCookieJar = None, type: str = None):
        """
        :param rcj: RequestsCookieJar对象，一般由应答头获得 response.cookies
        :param type: 需要更新到的cookie类型 DISPATCHER_TYPE.BUILD或DISPATCHER_TYPE.DEBUG
        :return: None
        """
        # 如果cookie池没有的话先创建默认值
        if session.get('dispatcher_trigger_type') == DISPATCHER_TRIGGER_TYPE.BY_SCHEDULE:
            project_id = session.get('project_id')
            if project_id:
                if project_id not in __class__.ScheduleCookiePool:
                    cls._add_empty_to_cookie_pool()
        else:
            session_id = session_id_manager.get_session_id()
            if session_id not in __class__.CookiePool:
                cls._add_empty_to_cookie_pool()
        # 如果入参rcj为空的RequestsCookieJar则不更新直接返回
        if len(rcj) == 0:
            return
        # 根据type和dispatcher_trigger_type进行cookie更新
        if type in [DISPATCHER_TYPE.DEBUG, DISPATCHER_TYPE.BUILD]:
            if session.get('dispatcher_trigger_type') == DISPATCHER_TRIGGER_TYPE.BY_SCHEDULE:
                project_id = session.get('project_id')
                if project_id:
                    old_rcj = __class__.ScheduleCookiePool[project_id]
                    if old_rcj is None:
                        __class__.ScheduleCookiePool[project_id] = rcj
                    else:
                        cls.update_request_cookie_jar(old_rcj=old_rcj, new_rcj=rcj)
            else:
                session_id = session_id_manager.get_session_id()
                old_rcj = __class__.CookiePool[session_id][type]
                if old_rcj is None:
                    __class__.CookiePool[session_id][type] = rcj
                else:
                    cls.update_request_cookie_jar(old_rcj=old_rcj, new_rcj=rcj)
        else:
            raise ValueError('不支持传入type值为%s，只支持type=single_case 或 type=build_case' % type)

    @classmethod
    def update_request_cookie_jar(cls, old_rcj: RequestsCookieJar, new_rcj: RequestsCookieJar):
        """将新的rcj更新到老的rcj上"""
        old_rcj.update(other=new_rcj)

    @classmethod
    def _add_empty_to_cookie_pool(cls):
        """
        为当前会话在cookie pool中添加一个新的，默认值为RequestsCookieJar()
        :return: None
        """
        if session.get('dispatcher_trigger_type') == DISPATCHER_TRIGGER_TYPE.BY_SCHEDULE:
            project_id = session.get('project_id')
            if project_id:
                __class__.ScheduleCookiePool.update({
                    project_id: RequestsCookieJar(),
                })
        else:
            session_id = session_id_manager.get_session_id()
            __class__.CookiePool.update({
                session_id: {
                    DISPATCHER_TYPE.DEBUG: RequestsCookieJar(),
                    DISPATCHER_TYPE.BUILD: RequestsCookieJar(),
                }
            })

    @classmethod
    def get_request_cookie_jar(cls, type: str) -> RequestsCookieJar:
        """
        获取当前会话中指定类型的cookie数据
        :param type: cookie类型 DISPATCHER_TYPE.BUILD或DISPATCHER_TYPE.DEBUG
        :return: RequestsCookieJar
        """
        if session.get('dispatcher_trigger_type') == DISPATCHER_TRIGGER_TYPE.BY_SCHEDULE:
            project_id = session.get('project_id')
            if project_id:
                if project_id not in __class__.ScheduleCookiePool:
                    cls._add_empty_to_cookie_pool()
                return __class__.ScheduleCookiePool[project_id]
        else:
            session_id = session_id_manager.get_session_id()
            if session_id not in __class__.CookiePool:
                cls._add_empty_to_cookie_pool()
            return __class__.CookiePool[session_id][type]

    @classmethod
    def clear_and_reset(cls, rcj: RequestsCookieJar, cookies: List[Mapping]):
        """
        清除并重置RequestsCookieJar
        :param rcj: 待重置的RequestsCookieJar，来自cookie pool中该会话cookie
        :param cookies: 新的cookie数据，是一个由字典组成的列表，每一个字典表示一个新的cookie
        :return: None
        """
        if rcj is None:
            rcj = RequestsCookieJar()
        rcj.clear()
        for cookie in cookies:
            rcj.set(**cookie)

    @classmethod
    def clear_cookie_pool(cls, type: str = None):
        """
        清空当前会话中的cookie数据
        :param type: cookie类型 DISPATCHER_TYPE.BUILD或DISPATCHER_TYPE.DEBUG
        :return:
        """
        if cls.get_request_cookie_jar(type=type) is not None:
            if type in [DISPATCHER_TYPE.DEBUG, DISPATCHER_TYPE.BUILD]:
                if session.get('dispatcher_trigger_type') == DISPATCHER_TRIGGER_TYPE.BY_SCHEDULE:
                    project_id = session.get('project_id')
                    if project_id:
                        __class__.ScheduleCookiePool[project_id] = RequestsCookieJar()
                else:
                    session_id = session_id_manager.get_session_id()
                    __class__.CookiePool[session_id][type] = RequestsCookieJar()
            else:
                raise ValueError('不支持传入type值为%s，只支持type=single_case 或 type=build_case' % type)
